import random
import numpy as np
import torch
from PIL import Image
from matplotlib import pyplot as plt
from torch.utils.data import Dataset
import PIL.ImageOps


# 自定义Dataset类，__getitem__(self,index)每次返回(img1, img2, 0/1)
class SiameseNetworkDataset(Dataset):

    def __init__(self, imageFolderDataset, transform=None, should_invert=True, gray=False):
        self.imageFolderDataset = imageFolderDataset
        self.transform = transform
        self.should_invert = should_invert
        self.gray = gray

    def __getitem__(self, index):
        # img0_tuple = random.choice(self.imageFolderDataset.imgs)  # 类别中任选一个
        img0_tuple = self.imageFolderDataset.imgs[index]
        should_get_same_class = random.randint(0, 1)  # 保证同类样本约占一半
        if should_get_same_class:
            while True:
                # 直到找到同一类别
                img1_tuple = random.choice(self.imageFolderDataset.imgs)
                if img0_tuple[1] == img1_tuple[1]:
                    break
        else:
            while True:
                # 直到找到非同一类别
                img1_tuple = random.choice(self.imageFolderDataset.imgs)
                if img0_tuple[1] != img1_tuple[1]:
                    break

        img0 = Image.open(img0_tuple[0])
        img1 = Image.open(img1_tuple[0])
        if self.gray:
            img0 = img0.convert("L")
            img1 = img1.convert("L")

        if self.should_invert:
            img0 = PIL.ImageOps.invert(img0)
            img1 = PIL.ImageOps.invert(img1)

        if self.transform is not None:
            img0 = self.transform(img0)
            img1 = self.transform(img1)

        return img0, img1, torch.from_numpy(np.array([int(img1_tuple[1] != img0_tuple[1])], dtype=np.float16))

    def __len__(self):
        return len(self.imageFolderDataset.imgs)


# 用于测试的Dataset类，__getitem__(self,index)每次返回(img, label)
class SiameseSingleDataset(Dataset):
    def __init__(self, imageFolderDataset, transform=None, should_invert=True,
                 gray=False):
        self.imageFolderDataset = imageFolderDataset
        self.transform = transform
        self.should_invert = should_invert
        self.gray = gray

    def __getitem__(self, index):
        img0_tuple = self.imageFolderDataset.imgs[index]

        img0 = Image.open(img0_tuple[0])
        if self.gray:
            img0 = img0.convert("L")
        if self.should_invert:
            img0 = PIL.ImageOps.invert(img0)
        # plt.imshow(img0)
        # plt.show()
        if self.transform is not None:
            img0 = self.transform(img0)
        it = torch.from_numpy(np.array(img0_tuple[1], dtype=np.long))

        return img0, it.long()

    # 這個 method 並不是 pytorch dataset 必要，只是方便未來我們想要指定「取哪幾張圖片」出來當作一個 batch 來 visualize
    def getbatch(self, indices):
        images = []
        labels = []
        for index in indices:
            image, label = self.__getitem__(index)
            images.append(image)
            labels.append(label)
        return torch.stack(images), torch.tensor(labels)

    def __len__(self):
        return len(self.imageFolderDataset.imgs)
